from math import sqrt

import pandas as pd
import numpy as np
import settings


class BackTest:

    def __init__(self):
        self.data_dict = {}  # 回测用数据
        self.data_frame = pd.DataFrame()  # 交易日志或单只股票价格走势
        self.trade_series = pd.Series()  # 第二天交易情况
        self.initial_money = settings.INITIAL_MONEY  # 初始资金
        self.final_money = self.initial_money  # 交易后资金

        self.max_money = 0  # 资产最高点
        self.min_money = 0  # 最高点后资产最低点

        self.roa = 0  # 资产回报率
        self.max_draw_down = 0  # 最大回撤率
        self.sharp_ratio = 0  # 夏普比率

    def run_back_test_with_stocks(self, data_dict):
        """
        多只股票日回测
        :param data_dict:
        :return:
        """
        self.data_dict = data_dict
        current_money = self.initial_money
        pre_share = {}  # 前一日持仓股数
        pre_money = current_money  # 前一天的资产
        self.max_money = current_money  # 资产最高点
        self.min_money = current_money  # 资产最低点

        date_list = sorted(self.data_dict.keys())
        stocks_name_list = self.data_dict[date_list[0]].index.values
        self.trade_series = pd.Series(False, index=stocks_name_list)

        # 日志
        self.data_frame['date_string'] = date_list
        self.data_frame['cumulative_return'] = ''  # 累计收益率
        self.data_frame['daily_return'] = ''  # 日收益率
        self.data_frame['daily_return_std'] = ''  # 波动率
        self.data_frame.set_index('date_string', drop=False, inplace=True)

        for stock_name in stocks_name_list:
            pre_share[stock_name] = 0

        for date_string in date_list:
            value_df = self.data_dict[date_string]

            value_df['share'] = ''  # 当日持仓股数
            value_df['equity'] = ''  # 持仓价值
            value_df.insert(0, 'is_holding', self.trade_series)
            value_df['is_holding'] = self.trade_series
            # holding_num = value_df['is_holding'].value_counts()[False]  # 当日计划持有股票数量
            holding_num = value_df['is_holding'].sum()  # 当日计划持有股票数量

            if holding_num > 0:
                capital = current_money / holding_num  # 每股资金，未取整股
                current_money = 0  # 不持有现金
            else:
                capital = 0

            for index, row in value_df.iterrows():
                if row['is_holding']:
                    row['share'] = capital / row['open']  # 当日持有份额
                    trade_duty = (pre_share[index] - row['share']) * row['open'] * settings.STAMP_DUTY
                    if trade_duty < 0:
                        trade_duty = 0  # 买入不收税
                    row['equity'] = row['share'] * row['close'] - trade_duty  # 持仓价值
                else:
                    row['share'] = 0
                    trade_duty = (pre_share[index] - row['share']) * row['open'] * settings.STAMP_DUTY
                    row['equity'] = - trade_duty

                pre_share[index] = row['share']
                value_df.loc[index] = row

            self.trade_series = value_df[settings.TRADE_FLAG]
            capital_after_close = value_df['equity'].sum()
            current_money = current_money + capital_after_close

            # 记录最大回撤
            if current_money > self.max_money:
                self.max_money = current_money
                self.min_money = current_money
            if current_money < self.min_money:
                self.min_money = current_money
                draw_down = (self.max_money - self.min_money) / self.max_money
                if draw_down > self.max_draw_down:
                    self.max_draw_down = draw_down

            # 绘图数据
            self.data_frame.loc[date_string, 'cumulative_return'] = \
                (current_money - self.initial_money) / self.initial_money * 100  # 总收益率
            self.data_frame.loc[date_string, 'daily_return'] = \
                (current_money - pre_money) / pre_money * 100

            self.data_frame.loc[date_string, 'daily_return_std'] = \
                round(np.std(self.data_frame.loc[date_string, 'daily_return']), 2)

        self.final_money = current_money
        self.roa = self.final_money / self.initial_money
        self.sharp_ratio = \
            ((self.data_frame['daily_return'].mean() - 0.00) / self.data_frame['daily_return'].std()) * \
            sqrt(settings.ANNUALIZATION_FACTORS['daily'])

    def run_back_test_with_stock(self, stock_data_frame):
        """
        单只股票回测
        :param stock_data_frame:
        :return:
        """
        self.data_frame = stock_data_frame
        is_holding = False
        current_money = self.initial_money
        for index, row in self.data_frame.iterrows():
            if not is_holding:  # 未持有状态
                if row[settings.TRADE_FLAG]:  # 买入
                    is_holding = True
            else:  # 持有状态
                current_money = current_money * (row['close'] / row['open'])
                if not row[settings.TRADE_FLAG]:
                    is_holding = False
        self.final_money = current_money
        self.roa = self.final_money / self.initial_money
